﻿//-----------------------------------------------------------------------
// <copyright file="Program.cs" company="Sergey Solyanik">
// Copyright (C) Sergey Solyanik.
//
// This file is subject to the terms and conditions of the Microsoft Public License (MS-PL).
// See http://www.microsoft.com/opensource/licenses.mspx#Ms-PL for more details.
// </copyright>
//----------------------------------------------------------------------- 
namespace Scriptster
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.IO;
    using System.Text;
    using System.Text.RegularExpressions;
    using System.Threading;

    using Microsoft.Win32;

    /// <summary>
    /// Various interpreter-related errors that we can emit.
    /// </summary>
    internal enum ErrorCode
    {
        /// <summary>
        /// Script file not found.
        /// </summary>
        ScriptNotFound = 1024,

        /// <summary>
        /// Compile failed.
        /// </summary>
        CompileFailed = 1025,

        /// <summary>
        /// Security exception - need to run this from an admin account.
        /// </summary>
        AccessDenied = 1026
    }

    /// <summary>
    /// The interpreter driver main class. Contains Main, does all the work.
    /// </summary>
    public class Program
    {
        /// <summary>
        /// Delegate type for async string reader.
        /// </summary>
        /// <returns> A read string. </returns>
        private delegate string StringDelegate();

        /// <summary>
        /// Processes the user input, takes the appropriate action (installs/uninstalls/compiles and runs the script).
        /// </summary>
        /// <param name="args"> Command line arguments.
        /// Expected to be: --install, --uninstall, --exec script [args] </param>
        public static void Main(string[] args)
        {
            string script = null;
            List<string> scriptArgs = null;
            bool verify = false;

            for (int i = 0; i < args.Length; ++i)
            {
                if (scriptArgs == null)
                {
                    if (args[i].Equals("--install"))
                    {
                        try
                        {
                            Install();
                        }
                        catch (System.Security.SecurityException)
                        {
                            Console.WriteLine("Error: Access Denied.");
                            Console.WriteLine("You need to run this application as an administrator to register.");
                            Environment.ExitCode = (int)ErrorCode.AccessDenied;
                        }

                        return;
                    }

                    if (args[i].Equals("--uninstall"))
                    {
                        try
                        {
                            UnInstall();
                        }
                        catch (System.Security.SecurityException)
                        {
                            Console.WriteLine("Error: Access Denied.");
                            Console.WriteLine("You need to run this application as an administrator to unregister.");
                            Environment.ExitCode = (int)ErrorCode.AccessDenied;
                        }

                        return;
                    }

                    if (args[i].Equals("--exec"))
                    {
                        if (i < args.Length - 1)
                        {
                            script = args[i + 1];
                            ++i;

                            scriptArgs = new List<string>();
                            continue;
                        }
                        else
                        {
                            Usage();
                            return;
                        }
                    }

                    if (args[i].Equals("--verify"))
                    {
                        if (i < args.Length - 1)
                        {
                            script = args[i + 1];
                            ++i;
                            verify = true;

                            continue;
                        }
                        else
                        {
                            Usage();
                            return;
                        }
                    }
                }

                if (scriptArgs == null)
                {
                    Console.WriteLine("Unrecognized command: {0}", args[i]);
                    Usage();
                    return;
                }

                scriptArgs.Add(args[i]);
            }

            if (scriptArgs == null && !verify)
            {
                Usage();
                return;
            }

            string fullScriptPath = Path.Combine(Environment.CurrentDirectory, script);

            if (!File.Exists(fullScriptPath))
            {
                Console.Error.WriteLine("{0} does not exist!", fullScriptPath);
                Environment.ExitCode = (int)ErrorCode.ScriptNotFound;
                return;
            }

            string tempDir = Path.Combine(Environment.GetEnvironmentVariable("TEMP"), "scriptster");
            Directory.CreateDirectory(tempDir);
            string exePathBase = Path.Combine(
                tempDir,
                fullScriptPath.Replace("_", "__").Replace(':', '_').Replace('\\', '_') + "_dir");
            string exePath = Path.Combine(exePathBase, "compiledscript.exe");

            if (verify ||
                (File.Exists(exePath) && File.GetLastWriteTime(exePath) < File.GetLastWriteTime(fullScriptPath)))
            {
                if (Directory.Exists(exePathBase))
                {
                    Directory.Delete(exePathBase, true);
                }
            }

            if (!File.Exists(exePath))
            {
                Directory.CreateDirectory(exePathBase);
                Compile(fullScriptPath, exePath, exePathBase, verify);
            }

            if (!File.Exists(exePath))
            {
                Console.Error.WriteLine("Could not compile {0}!", fullScriptPath);
                Environment.ExitCode = (int)ErrorCode.CompileFailed;
                return;
            }

            if (verify)
            {
                return;
            }

            StringBuilder sb = new StringBuilder();
            foreach (string a in scriptArgs)
            {
                if (sb.Length > 0)
                {
                    sb.Append(' ');
                }

                sb.Append('"');
                sb.Append(a);
                sb.Append('"');
            }

            Process compiledScript = new Process();
            compiledScript.StartInfo.UseShellExecute = false;
            compiledScript.StartInfo.FileName = exePath;
            compiledScript.StartInfo.Arguments = sb.ToString();

            compiledScript.Start();
            compiledScript.WaitForExit();

            Environment.ExitCode = compiledScript.ExitCode;
        }

        /// <summary>
        /// Prints out usage statement.
        /// </summary>
        private static void Usage()
        {
            Console.WriteLine("scriptster --install");
            Console.WriteLine("scriptster --uninstall");
            Console.WriteLine("scriptster --exec scriptfile arg1...argN");
            Console.WriteLine("scriptster --verify scriptfile");
        }

        /// <summary>
        /// Installs C# scripter.
        /// </summary>
        private static void Install()
        {
            string pathExt = Environment.GetEnvironmentVariable("PATHEXT", System.EnvironmentVariableTarget.Machine);
            if (!pathExt.ToUpper().Contains(".CSSCRIPT"))
            {
                Environment.SetEnvironmentVariable("PATHEXT", ".CSSCRIPT;" + pathExt,
                    EnvironmentVariableTarget.Machine);
            }

            RegistryKey key = Registry.ClassesRoot.CreateSubKey(".csscript");
            key.SetValue(null, "scriptsterfile");
            key.Close();
            key = Registry.ClassesRoot.CreateSubKey("scriptsterfile");
            key.Close();
            key = Registry.ClassesRoot.CreateSubKey("scriptsterfile\\shell");
            key.Close();
            key = Registry.ClassesRoot.CreateSubKey("scriptsterfile\\shell\\open");
            key.Close();
            key = Registry.ClassesRoot.CreateSubKey("scriptsterfile\\shell\\open\\command");
            key.SetValue(
                null,
                "\"" + System.Reflection.Assembly.GetExecutingAssembly().Location + "\" --exec \"%1\" %*");
            key.Close();

            Console.WriteLine("Please restart the command prompt to complete the installation.");
        }

        /// <summary>
        /// Uninstalls C# scripter.
        /// </summary>
        private static void UnInstall()
        {
            string pathExt = Environment.GetEnvironmentVariable("PATHEXT", System.EnvironmentVariableTarget.Machine);
            if (pathExt.Contains(".CSSCRIPT"))
            {
                string newPathExt = pathExt.Replace(".CSSCRIPT;", String.Empty);
                newPathExt = newPathExt.Replace(";.CSSCRIPT", String.Empty);
                newPathExt = newPathExt.Replace(".CSSCRIPT", String.Empty);

                Environment.SetEnvironmentVariable(
                    "PATHEXT",
                    newPathExt,
                    EnvironmentVariableTarget.Machine);
            }

            Registry.ClassesRoot.DeleteSubKey(".csscript", false);
            try
            {
                Registry.ClassesRoot.DeleteSubKeyTree("scriptsterfile");
            }
            catch (ArgumentException)
            {
                // This is thrown if the key is missing. Eat it - it's OK.
            }

            string tempDir = Path.Combine(Environment.GetEnvironmentVariable("TEMP"), "scriptster");
            try
            {
                Directory.Delete(tempDir, true);
            }
            catch (DirectoryNotFoundException)
            {
            }
        }

        /// <summary>
        /// Compiles a script.
        /// </summary>
        /// <param name="scriptName"> Name of the script. </param>
        /// <param name="outputName"> Name of the output file. </param>
        /// <param name="outputDir"> Directory where all the dependencies go. </param>
        /// <param name="verbose"> If true, compiler output is printed out. </param>
        private static void Compile(string scriptName, string outputName, string outputDir, bool verbose)
        {
            string runtimeEnvironmentPath = System.Runtime.InteropServices.RuntimeEnvironment.GetRuntimeDirectory();
            if (runtimeEnvironmentPath.EndsWith("\\"))
            {
                runtimeEnvironmentPath = runtimeEnvironmentPath.Substring(0, runtimeEnvironmentPath.Length - 1);
            }

            string compilerExe = Path.Combine(runtimeEnvironmentPath, "csc.exe");
            bool isFramework35 = false;

            // However, if it's 2.0, there is a compiler-only extension (3.5). If 3.5 is indeed installed, use that.
            string baseFrameworkPath = Path.GetDirectoryName(runtimeEnvironmentPath);
            string runtimeName = Path.GetFileName(runtimeEnvironmentPath);
            if ("v2.0.50727".Equals(runtimeName))
            {
                string compilerExe35 = Path.Combine(Path.Combine(baseFrameworkPath, "v3.5"), "csc.exe");
                if (File.Exists(compilerExe35))
                {
                    compilerExe = compilerExe35;
                    isFramework35 = true;
                }
            }

            List<string> references = new List<string>();
            references.Add("System.dll");
            references.Add("System.Core.dll");
            references.Add("System.Data.dll");
            references.Add("System.Xml.dll");
            if (isFramework35)
            {
                references.Add("System.Data.Linq.dll");
            }

            string codeFile = Path.Combine(outputDir, "script.cs");
            if (!PreprocessFile(scriptName, codeFile, references))
            {
                return;
            }

            CopyReferences(references, outputDir);

            if (isFramework35)
            {
                references = ReprocessReferences(references);
            }

            StringBuilder sb = new StringBuilder();
            sb.Append("/target:exe /debug- /optimize+ /checked- /unsafe+ /noconfig ");

            foreach (string r in references)
            {
                sb.Append("/reference:\"" + r + "\" ");
            }

            sb.Append("/out:\"" + outputName + "\" ");
            sb.Append("\"" + codeFile + "\"");

            string args = sb.ToString();

            if (verbose)
            {
                Console.WriteLine("{0} {1}", compilerExe, args);
            }

            Process compiler = new Process();
            compiler.StartInfo.UseShellExecute = false;
            compiler.StartInfo.RedirectStandardError = true;
            compiler.StartInfo.RedirectStandardOutput = true;
            compiler.StartInfo.CreateNoWindow = true;
            compiler.StartInfo.FileName = compilerExe;
            compiler.StartInfo.Arguments = args;

            compiler.Start();

            StringDelegate outputStreamAsyncReader = new StringDelegate(compiler.StandardOutput.ReadToEnd);
            StringDelegate errorStreamAsyncReader = new StringDelegate(compiler.StandardError.ReadToEnd);
            IAsyncResult outAsyncResult = outputStreamAsyncReader.BeginInvoke(null, null);
            IAsyncResult errAsyncResult = errorStreamAsyncReader.BeginInvoke(null, null);

            WaitHandle[] handles = { outAsyncResult.AsyncWaitHandle, errAsyncResult.AsyncWaitHandle };
            if (!WaitHandle.WaitAll(handles))
            {
                Console.WriteLine("Execution aborted!");
                return;
            }

            string results = outputStreamAsyncReader.EndInvoke(outAsyncResult);
            string errorMessage = errorStreamAsyncReader.EndInvoke(errAsyncResult);

            compiler.WaitForExit();

            if (!String.IsNullOrEmpty(errorMessage))
            {
                Console.WriteLine("{0}", errorMessage);
            }

            if (!String.IsNullOrEmpty(results))
            {
                string scriptProperName = Path.GetFileName(scriptName);
                string codeName = Path.GetFileName(codeFile);

                StringReader sr = new StringReader(results);
                int skipLines = verbose ? 0 : 4;
                string line;
                while ((line = sr.ReadLine()) != null)
                {
                    if (--skipLines < 0)
                    {
                        Console.WriteLine(line.Replace(codeName, scriptProperName));
                    }
                }
            }
        }

        /// <summary>
        /// Preprocesses the script file.
        /// </summary>
        /// <param name="scriptName"> The script file name. </param>
        /// <param name="codeFile"> The name of the file where C# code will go. </param>
        /// <param name="references"> A list of assembly references. </param>
        /// <returns> True if the file was preprocessed successfully. </returns>
        private static bool PreprocessFile(string scriptName, string codeFile, List<string> references)
        {
            bool preprocessorError = false;
            string scriptProperName = Path.GetFileName(scriptName);

            Regex referenceMatcher = new Regex("^//#ref\\s+(\".+\"|\\S+)\\s*$");

            StreamReader reader = new StreamReader(scriptName);
            StreamWriter writer = new StreamWriter(codeFile);
            string line;
            int lineNumber = 1;
            while ((line = reader.ReadLine()) != null)
            {
                Match m = referenceMatcher.Match(line);
                if (m.Success)
                {
                    string reference = m.Groups[1].Value;
                    reference = reference.Replace("\"", String.Empty);
                    reference = Environment.ExpandEnvironmentVariables(reference);
                    references.Add(reference);
                    writer.WriteLine();
                }
                else
                {
                    writer.WriteLine(line);
                }

                ++lineNumber;
            }

            reader.Close();
            writer.Close();

            return !preprocessorError;
        }

        /// <summary>
        /// Takes a list of system default references, and matches them to a version of a framework.
        /// </summary>
        /// <param name="references"> A list of references. </param>
        /// <returns> A list of references where 3.5 dlls are explicitly qualified with path. </returns>
        private static List<string> ReprocessReferences(List<string> references)
        {
            string referencePath = Path.Combine(
                Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles),
                "Reference Assemblies\\Microsoft\\Framework\\v3.5");

            List<string> result = new List<string>();
            foreach (string reference in references)
            {
                string substitute = Path.Combine(referencePath, reference);
                if (File.Exists(substitute))
                {
                    result.Add(substitute);
                }
                else
                {
                    result.Add(reference);
                }
            }

            return result;
        }

        /// <summary>
        /// Copies non-GAC'ed references to the output directory.
        /// </summary>
        /// <param name="references"> The list of references. </param>
        /// <param name="outputDir"> Target directory. </param>
        private static void CopyReferences(List<string> references, string outputDir)
        {
            foreach (string reference in references)
            {
                string sourceDir = Path.GetDirectoryName(reference);
                if (!String.IsNullOrEmpty(sourceDir) && File.Exists(reference))
                {
                    // If we have a full path at this point, copy the reference
                    string targetFile = Path.Combine(outputDir, Path.GetFileName(reference));
                    File.Copy(reference, targetFile, true);
                    FileAttributes cleanAttributes = File.GetAttributes(targetFile) &
                        ~(FileAttributes.System | FileAttributes.Hidden | FileAttributes.ReadOnly);
                    File.SetAttributes(targetFile, cleanAttributes);
                }
            }
        }
    }
}
